{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# 增强RAG的关联段落提取（RSE）\n",
    "\n",
    "实现关联段落提取（Relevant Segment Extraction，RSE）技术，以提高RAG系统的上下文质量。我们不仅仅检索一组孤立的片段，而是识别并重建提供更好上下文的连续文本段落，从而为语言模型提供更好的支持。\n",
    "\n",
    "-----\n",
    "核心概念：\n",
    "\n",
    "相关的片段往往会在文档中聚集成簇。通过识别这些簇并保持其连续性，RSE为大型语言模型提供了更加连贯的上下文。\n",
    "\n",
    "-----\n",
    "实现步骤：\n",
    "- 处理文档以创建向量存储：从PDF 中提取文本，分割文本块（0重叠）并创建向量存储\n",
    "- 根据查询计算相关性分数和块值：\n",
    "    - 先获取所有带有相似度分数的块；\n",
    "    - 获取相关性分数，如果不在结果中则默认为0，同时应用惩罚以将不相关的块转换为负值\n",
    "- 根据块值找到最佳文本块：使用最大子数组和算法的变体找到最佳\n",
    "- 从最佳块中重建文本段落：基于块索引重建文本段落\n",
    "- 利用将段落格式化为上下文生成回答"
   ],
   "id": "ca11d6758ecdaf1d"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-24T11:22:05.233442Z",
     "start_time": "2025-04-24T11:22:02.668534Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import fitz\n",
    "import os\n",
    "import re\n",
    "import json\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "from openai import OpenAI\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()"
   ],
   "id": "69a04ea36e6e795b",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-24T11:22:06.964460Z",
     "start_time": "2025-04-24T11:22:06.623435Z"
    }
   },
   "cell_type": "code",
   "source": [
    "client = OpenAI(\n",
    "    base_url=os.getenv(\"LLM_BASE_URL\"),\n",
    "    api_key=os.getenv(\"LLM_API_KEY\")\n",
    ")\n",
    "llm_model = os.getenv(\"LLM_MODEL_ID\")\n",
    "embedding_model = os.getenv(\"EMBEDDING_MODEL_ID\")\n",
    "\n",
    "pdf_path = \"../../data/AI_Information.en.zh-CN.pdf\""
   ],
   "id": "6a81c5a9ed9e57aa",
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-24T11:22:08.636910Z",
     "start_time": "2025-04-24T11:22:08.632962Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def extract_text_from_pdf(pdf_path):\n",
    "    \"\"\"\n",
    "    从 PDF 文件中提取文本，并打印前 `num_chars` 个字符。\n",
    "\n",
    "    Args:\n",
    "    pdf_path (str): Path to the PDF file.\n",
    "\n",
    "    Returns:\n",
    "    str: Extracted text from the PDF.\n",
    "    \"\"\"\n",
    "    # 打开 PDF 文件\n",
    "    mypdf = fitz.open(pdf_path)\n",
    "    all_text = \"\"  # 初始化一个空字符串以存储提取的文本\n",
    "\n",
    "    # Iterate through each page in the PDF\n",
    "    for page_num in range(mypdf.page_count):\n",
    "        page = mypdf[page_num]\n",
    "        text = page.get_text(\"text\")  # 从页面中提取文本\n",
    "        all_text += text  # 将提取的文本追加到 all_text 字符串中\n",
    "\n",
    "    return all_text  # 返回提取的文本"
   ],
   "id": "f5c8e552170e7376",
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-24T11:22:09.885186Z",
     "start_time": "2025-04-24T11:22:09.881515Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def chunk_text(text, n, overlap):\n",
    "    \"\"\"\n",
    "    将文本分割为重叠的块\n",
    "\n",
    "    Args:\n",
    "        text (str): 要分割的文本\n",
    "        n (int): 每个块的字符数\n",
    "        overlap (int): 块之间的重叠字符数\n",
    "\n",
    "    Returns:\n",
    "        List[str]: 文本块列表\n",
    "    \"\"\"\n",
    "    chunks = []  #\n",
    "    for i in range(0, len(text), n - overlap):\n",
    "        # 添加从当前索引到索引 + 块大小的文本块\n",
    "        chunk = text[i:i + n]\n",
    "        if chunk:\n",
    "            chunks.append(chunk)\n",
    "\n",
    "    return chunks  # Return the list of text chunks"
   ],
   "id": "1ab9d23cb4fdfa7b",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-24T11:22:12.582385Z",
     "start_time": "2025-04-24T11:22:12.576199Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class SimpleVectorStore:\n",
    "    \"\"\"\n",
    "    使用 NumPy 实现的轻量级向量存储。\n",
    "    \"\"\"\n",
    "    def __init__(self, dimension=1536):\n",
    "        \"\"\"\n",
    "        初始化向量存储。\n",
    "\n",
    "        Args:\n",
    "            dimension (int): 嵌入向量的维度\n",
    "        \"\"\"\n",
    "        self.dimension = dimension  # 嵌入向量的维度\n",
    "        self.vectors = []  # 存储嵌入向量的列表\n",
    "        self.documents = []  # 存储文档片段的列表\n",
    "        self.metadata = []  # 存储元数据的列表\n",
    "\n",
    "    def add_documents(self, documents, vectors=None, metadata=None):\n",
    "        \"\"\"\n",
    "        向向量存储中添加文档。\n",
    "\n",
    "        Args:\n",
    "            documents (List[str]): 文档片段列表\n",
    "            vectors (List[List[float]], 可选): 嵌入向量列表\n",
    "            metadata (List[Dict], 可选): 元数据字典列表\n",
    "        \"\"\"\n",
    "        if vectors is None:  # 如果未提供向量，则生成一个空列表\n",
    "            vectors = [None] * len(documents)\n",
    "\n",
    "        if metadata is None:  # 如果未提供元数据，则生成一个空字典列表\n",
    "            metadata = [{} for _ in range(len(documents))]\n",
    "\n",
    "        for doc, vec, meta in zip(documents, vectors, metadata):  # 遍历文档、向量和元数据\n",
    "            self.documents.append(doc)  # 将文档片段添加到列表中\n",
    "            self.vectors.append(vec)  # 将嵌入向量添加到列表中\n",
    "            self.metadata.append(meta)  # 将元数据添加到列表中\n",
    "\n",
    "    def search(self, query_vector, top_k=5):\n",
    "        \"\"\"\n",
    "        搜索最相似的文档。\n",
    "\n",
    "        Args:\n",
    "            query_vector (List[float]): 查询嵌入向量\n",
    "            top_k (int): 返回的结果数量\n",
    "\n",
    "        Returns:\n",
    "            List[Dict]: 包含文档、分数和元数据的结果列表\n",
    "        \"\"\"\n",
    "        if not self.vectors or not self.documents:  # 如果向量或文档为空，返回空列表\n",
    "            return []\n",
    "\n",
    "        # 将查询向量转换为 NumPy 数组\n",
    "        query_array = np.array(query_vector)\n",
    "\n",
    "        # 计算相似度\n",
    "        similarities = []\n",
    "        for i, vector in enumerate(self.vectors):  # 遍历存储中的向量\n",
    "            if vector is not None:  # 如果向量不为空\n",
    "                # 计算余弦相似度\n",
    "                similarity = np.dot(query_array, vector) / (\n",
    "                    np.linalg.norm(query_array) * np.linalg.norm(vector)\n",
    "                )\n",
    "                similarities.append((i, similarity))  # 将索引和相似度添加到列表中\n",
    "\n",
    "        # 按相似度降序排序\n",
    "        similarities.sort(key=lambda x: x[1], reverse=True)\n",
    "\n",
    "        # 获取前 top-k 结果\n",
    "        results = []\n",
    "        for i, score in similarities[:top_k]:  # 遍历前 top-k 的结果\n",
    "            results.append({\n",
    "                \"document\": self.documents[i],  # 文档片段\n",
    "                \"score\": float(score),  # 相似度分数\n",
    "                \"metadata\": self.metadata[i]  # 元数据\n",
    "            })\n",
    "\n",
    "        return results  # 返回结果列表\n"
   ],
   "id": "f28a9a9558ab068b",
   "outputs": [],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-24T11:22:15.506301Z",
     "start_time": "2025-04-24T11:22:15.502552Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def create_embeddings(texts):\n",
    "    \"\"\"\n",
    "    为文本生成嵌入向量。\n",
    "\n",
    "    Args:\n",
    "        texts (List[str]): 要嵌入的文本列表\n",
    "\n",
    "    Returns:\n",
    "        List[List[float]]: 嵌入向量列表\n",
    "    \"\"\"\n",
    "    if not texts:  # 如果没有提供文本，返回空列表\n",
    "        return []\n",
    "\n",
    "    # 如果列表很长，则按批次处理\n",
    "    batch_size = 100  # 根据API限制进行调整\n",
    "    all_embeddings = []  # 初始化一个列表来存储所有嵌入向量\n",
    "\n",
    "    for i in range(0, len(texts), batch_size):  # 按批次处理文本\n",
    "        batch = texts[i:i + batch_size]  # 获取当前批次的文本\n",
    "\n",
    "        # 使用指定的模型为当前批次生成嵌入向量\n",
    "        response = client.embeddings.create(\n",
    "            input=batch,  # 输入文本批次\n",
    "            model=embedding_model  # 使用的模型\n",
    "        )\n",
    "\n",
    "        # 从响应中提取嵌入向量\n",
    "        batch_embeddings = [item.embedding for item in response.data]\n",
    "        all_embeddings.extend(batch_embeddings)  # 将批次嵌入向量添加到总列表中\n",
    "\n",
    "    return all_embeddings  # 返回所有嵌入向量的列表\n"
   ],
   "id": "a7991ed0c43a2a11",
   "outputs": [],
   "execution_count": 6
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 用 RSE 处理文档",
   "id": "12caa71d4b9a99e3"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-24T11:22:18.486637Z",
     "start_time": "2025-04-24T11:22:18.482541Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def process_document(pdf_path, chunk_size=800):\n",
    "    \"\"\"\n",
    "    处理文档以供 RSE 使用。\n",
    "\n",
    "    Args:\n",
    "        pdf_path (str): PDF 文档的路径\n",
    "        chunk_size (int): 每个块的字符大小\n",
    "\n",
    "    Returns:\n",
    "        Tuple[List[str], SimpleVectorStore, Dict]: 块列表、向量存储和文档信息\n",
    "    \"\"\"\n",
    "    print(\"从文档中提取文本...\")\n",
    "    # 从 PDF 文件中提取文本\n",
    "    text = extract_text_from_pdf(pdf_path)\n",
    "\n",
    "    print(\"将文本切分为非重叠段落...\")\n",
    "    # 将提取的文本切分为非重叠段落\n",
    "    chunks = chunk_text(text, n=chunk_size, overlap=0)\n",
    "    print(f\"创建了 {len(chunks)} 个块\")\n",
    "\n",
    "    print(\"为块生成嵌入向量...\")\n",
    "    # 为文本块生成嵌入向量\n",
    "    chunk_embeddings = create_embeddings(chunks)\n",
    "\n",
    "    # 创建 SimpleVectorStore 的实例\n",
    "    vector_store = SimpleVectorStore()\n",
    "\n",
    "    # 添加带有元数据的文档（包括块索引以便后续重建）\n",
    "    metadata = [{\"chunk_index\": i, \"source\": pdf_path} for i in range(len(chunks))]\n",
    "    vector_store.add_documents(chunks, chunk_embeddings, metadata)\n",
    "\n",
    "    # 跟踪原始文档结构以便段落重建\n",
    "    doc_info = {\n",
    "        \"chunks\": chunks,\n",
    "        \"source\": pdf_path,\n",
    "    }\n",
    "\n",
    "    return chunks, vector_store, doc_info\n"
   ],
   "id": "cbecedb04ac8fefd",
   "outputs": [],
   "execution_count": 7
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## RSE核心算法：计算块值和寻找最佳段落\n",
    "现在我们已经有了处理文档和为其块生成嵌入向量所需的函数，可以实现RSE的核心算法。"
   ],
   "id": "b4694934c8ed71f2"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-24T11:22:20.650838Z",
     "start_time": "2025-04-24T11:22:20.646942Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def calculate_chunk_values(query, chunks, vector_store, irrelevant_chunk_penalty=0.2):\n",
    "    \"\"\"\n",
    "    通过结合相关性和位置计算块的值。\n",
    "\n",
    "    Args:\n",
    "        query (str): 查询文本\n",
    "        chunks (List[str]): 文档块列表\n",
    "        vector_store (SimpleVectorStore): 包含块的向量存储\n",
    "        irrelevant_chunk_penalty (float): 不相关块的惩罚值\n",
    "\n",
    "    Returns:\n",
    "        List[float]: 块值列表\n",
    "    \"\"\"\n",
    "    # 创建查询嵌入\n",
    "    query_embedding = create_embeddings([query])[0]\n",
    "\n",
    "    # 获取所有带有相似度分数的块\n",
    "    num_chunks = len(chunks)\n",
    "    results = vector_store.search(query_embedding, top_k=num_chunks)\n",
    "\n",
    "    # 创建从块索引到相关性分数的映射\n",
    "    relevance_scores = {result[\"metadata\"][\"chunk_index\"]: result[\"score\"] for result in results}\n",
    "\n",
    "    # 计算块值（相关性分数减去惩罚）\n",
    "    chunk_values = []\n",
    "    for i in range(num_chunks):\n",
    "        # 获取相关性分数，如果不在结果中则默认为0\n",
    "        score = relevance_scores.get(i, 0.0)\n",
    "        # 应用惩罚以将不相关的块转换为负值\n",
    "        value = score - irrelevant_chunk_penalty\n",
    "        chunk_values.append(value)\n",
    "\n",
    "    return chunk_values\n"
   ],
   "id": "17acb37144675e6c",
   "outputs": [],
   "execution_count": 8
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-24T11:22:22.339225Z",
     "start_time": "2025-04-24T11:22:22.333907Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def find_best_segments(chunk_values, max_segment_length=20, total_max_length=30, min_segment_value=0.2):\n",
    "    \"\"\"\n",
    "    使用最大子数组和算法的变体找到最佳段落。\n",
    "\n",
    "    Args:\n",
    "        chunk_values (List[float]): 每个块的值\n",
    "        max_segment_length (int): 单个段落的最大长度\n",
    "        total_max_length (int): 所有段落的最大总长度\n",
    "        min_segment_value (float): 被考虑的段落的最小值\n",
    "\n",
    "    Returns:\n",
    "        List[Tuple[int, int]]: 最佳段落的（开始，结束）索引列表\n",
    "    \"\"\"\n",
    "    print(\"寻找最佳连续文本段落...\")\n",
    "\n",
    "    best_segments = []\n",
    "    segment_scores = []\n",
    "    total_included_chunks = 0\n",
    "\n",
    "    # 继续寻找段落直到达到限制\n",
    "    while total_included_chunks < total_max_length:\n",
    "        best_score = min_segment_value  # 段落的最低阈值\n",
    "        best_segment = None\n",
    "\n",
    "        # 尝试每个可能的起始位置\n",
    "        for start in range(len(chunk_values)):\n",
    "            # 如果该起始位置已经在选定的段落中，则跳过(重叠内容部分)\n",
    "            if any(start >= s[0] and start < s[1] for s in best_segments):\n",
    "                continue\n",
    "\n",
    "            # 尝试每个可能的段落长度\n",
    "            for length in range(1, min(max_segment_length, len(chunk_values) - start) + 1):\n",
    "                end = start + length\n",
    "\n",
    "                # 如果结束位置已经在选定的段落中，则跳过\n",
    "                if any(end > s[0] and end <= s[1] for s in best_segments):\n",
    "                    continue\n",
    "\n",
    "                # 计算段落值为块值的总和\n",
    "                segment_value = sum(chunk_values[start:end])\n",
    "\n",
    "                # 如果这个段落更好，则更新最佳段落\n",
    "                if segment_value > best_score:\n",
    "                    best_score = segment_value\n",
    "                    best_segment = (start, end)\n",
    "\n",
    "        # 如果找到了一个好的段落，则添加它\n",
    "        if best_segment:\n",
    "            best_segments.append(best_segment)\n",
    "            segment_scores.append(best_score)\n",
    "            total_included_chunks += best_segment[1] - best_segment[0]\n",
    "            print(f\"找到段落 {best_segment}，得分 {best_score:.4f}\")\n",
    "        else:\n",
    "            # 没有更多的好段落可找\n",
    "            break\n",
    "\n",
    "    # 按段落的起始位置排序以便于阅读\n",
    "    best_segments = sorted(best_segments, key=lambda x: x[0])\n",
    "\n",
    "    return best_segments, segment_scores\n"
   ],
   "id": "a0f16f4d36cbdd99",
   "outputs": [],
   "execution_count": 9
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## 为RAG重建和使用段落\n",
    "\n",
    "通过重建和利用文本段落来改进上下文质量和生成效果。"
   ],
   "id": "f5f12d824811bc79"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-24T11:22:25.632514Z",
     "start_time": "2025-04-24T11:22:25.628489Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def reconstruct_segments(chunks, best_segments):\n",
    "    \"\"\"\n",
    "    基于块索引重建文本段落。\n",
    "\n",
    "    Args:\n",
    "        chunks (List[str]): 所有文档块的列表\n",
    "        best_segments (List[Tuple[int, int]]): 段落的（开始，结束）索引列表\n",
    "\n",
    "    Returns:\n",
    "        List[Dict]: 重建的文本段落列表，每个段落包含文本和其范围\n",
    "    \"\"\"\n",
    "    reconstructed_segments = []  # 初始化一个空列表以存储重建的段落\n",
    "\n",
    "    for start, end in best_segments:\n",
    "        # 将此段落中的块连接起来以形成完整的段落文本\n",
    "        segment_text = \" \".join(chunks[start:end])\n",
    "        # 将段落文本及其范围追加到重建的段落列表中\n",
    "        reconstructed_segments.append({\n",
    "            \"text\": segment_text,  # 段落文本\n",
    "            \"segment_range\": (start, end),  # 段落范围\n",
    "        })\n",
    "\n",
    "    return reconstructed_segments  # 返回重建的文本段落列表\n"
   ],
   "id": "7ea5390833d6606c",
   "outputs": [],
   "execution_count": 10
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-24T11:22:27.815551Z",
     "start_time": "2025-04-24T11:22:27.811711Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def format_segments_for_context(segments):\n",
    "    \"\"\"\n",
    "    将段落格式化为适用于LLM的上下文字符串。\n",
    "\n",
    "    Args:\n",
    "        segments (List[Dict]): 段落字典列表\n",
    "\n",
    "    Returns:\n",
    "        str: 格式化后的上下文文本\n",
    "    \"\"\"\n",
    "    context = []  # 初始化一个空列表以存储格式化后的上下文\n",
    "\n",
    "    for i, segment in enumerate(segments):\n",
    "        # 为每个段落创建一个包含索引和块范围的标题\n",
    "        # segment_header = f\"SEGMENT {i+1} (Chunks {segment['segment_range'][0]}-{segment['segment_range'][1]-1}):\"\n",
    "        segment_header = f\"分段{i+1}（包含文本块{segment['segment_range'][0]}至{segment['segment_range'][1]-1}）：\"\n",
    "        context.append(segment_header)  # 将段落标题添加到上下文列表中\n",
    "        context.append(segment['text'])  # 将段落文本添加到上下文列表中\n",
    "        context.append(\"-\" * 80)  # 添加分隔线以提高可读性\n",
    "\n",
    "    # 将上下文列表中的所有元素用双换行符连接并返回结果\n",
    "    return \"\\n\\n\".join(context)\n"
   ],
   "id": "29185f8ac4b3fe42",
   "outputs": [],
   "execution_count": 11
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 使用 RSE 上下文生成回答",
   "id": "9a370db2723f29b7"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-24T11:22:30.833252Z",
     "start_time": "2025-04-24T11:22:30.827293Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def generate_response(query, context):\n",
    "    \"\"\"\n",
    "    根据查询和上下文生成响应。\n",
    "\n",
    "    Args:\n",
    "        query (str): 用户查询\n",
    "        context (str): 来自相关段落的上下文文本\n",
    "\n",
    "    Returns:\n",
    "        str: 生成的响应\n",
    "    \"\"\"\n",
    "    print(\"正在使用相关段落作为上下文生成响应...\")\n",
    "\n",
    "    # 定义系统提示以引导AI的行为\n",
    "    system_prompt = \"\"\"\n",
    "    您是基于上下文智能应答的AI助手，需根据提供的文档段落回答用户问题。\n",
    "    这些文档段落是通过相关性检索匹配到当前问题的上下文内容。\n",
    "    请严格依据以下要求执行：\n",
    "    1. 整合分析所有相关段落信息\n",
    "    2. 生成全面准确的综合回答\n",
    "    3. 当上下文不包含有效信息时，必须明确告知无法回答\n",
    "    \"\"\"\n",
    "\n",
    "    # 通过组合上下文和查询创建用户提示\n",
    "    user_prompt = f\"\"\"\n",
    "    上下文内容：\n",
    "    {context}\n",
    "\n",
    "    问题：{query}\n",
    "\n",
    "    请基于上述上下文内容提供专业可靠的回答。\n",
    "    \"\"\"\n",
    "\n",
    "    # 使用指定的模型生成响应\n",
    "    response = client.chat.completions.create(\n",
    "        model=llm_model,\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": system_prompt},\n",
    "            {\"role\": \"user\", \"content\": user_prompt}\n",
    "        ],\n",
    "        temperature=0\n",
    "    )\n",
    "\n",
    "    # 返回生成的响应内容\n",
    "    return response.choices[0].message.content\n"
   ],
   "id": "46185fac6f91d395",
   "outputs": [],
   "execution_count": 12
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 完整的RSE管道功能",
   "id": "56067dd7e8b08ffa"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-24T11:22:36.652518Z",
     "start_time": "2025-04-24T11:22:36.648788Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def rag_with_rse(pdf_path, query, chunk_size=800, irrelevant_chunk_penalty=0.2):\n",
    "    \"\"\"\n",
    "    完整的RAG管道，包含相关段落提取（Relevant Segment Extraction）。\n",
    "\n",
    "    Args:\n",
    "        pdf_path (str): 文档路径\n",
    "        query (str): 用户查询\n",
    "        chunk_size (int): 每个块的大小\n",
    "        irrelevant_chunk_penalty (float): 对不相关块的惩罚值\n",
    "\n",
    "    Returns:\n",
    "        Dict: 包含查询、段落和响应的结果字典\n",
    "    \"\"\"\n",
    "    print(\"\\n=== 开始带有相关段落提取的RAG ===\")\n",
    "    print(f\"查询: {query}\")\n",
    "\n",
    "    # 处理文档以提取文本、分块并创建嵌入\n",
    "    chunks, vector_store, doc_info = process_document(pdf_path, chunk_size)\n",
    "\n",
    "    # 根据查询计算相关性分数和块值\n",
    "    print(\"\\n计算相关性分数和块值...\")\n",
    "    chunk_values = calculate_chunk_values(query, chunks, vector_store, irrelevant_chunk_penalty)\n",
    "\n",
    "    # 根据块值找到最佳文本段落\n",
    "    best_segments, scores = find_best_segments(\n",
    "        chunk_values,\n",
    "        max_segment_length=20,  # 最大段落长度\n",
    "        total_max_length=30,  # 所有段落的最大总长度\n",
    "        min_segment_value=0.2  # 考虑段落的最小值\n",
    "    )\n",
    "\n",
    "    # 从最佳块中重建文本段落\n",
    "    print(\"\\n从块中重建文本段落...\")\n",
    "    segments = reconstruct_segments(chunks, best_segments)\n",
    "\n",
    "    # 将段落格式化为语言模型的上下文字符串\n",
    "    context = format_segments_for_context(segments)\n",
    "\n",
    "    # 使用上下文从语言模型生成响应\n",
    "    response = generate_response(query, context)\n",
    "\n",
    "    # 将结果编译成字典\n",
    "    result = {\n",
    "        \"query\": query,  # 用户查询\n",
    "        \"segments\": segments,  # 提取的段落\n",
    "        \"response\": response  # 模型生成的响应\n",
    "    }\n",
    "\n",
    "    print(\"\\n=== 最终响应 ===\")\n",
    "    print(response)\n",
    "\n",
    "    return result\n"
   ],
   "id": "eb274f870768fa3f",
   "outputs": [],
   "execution_count": 13
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## 与标准检索进行比较\n",
    "实现一种标准的检索方法，以便与RSE（相关段落提取）进行比较："
   ],
   "id": "d31f6bc347007e72"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-24T11:22:38.871163Z",
     "start_time": "2025-04-24T11:22:38.865068Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def standard_top_k_retrieval(pdf_path, query, k=10, chunk_size=800):\n",
    "    \"\"\"\n",
    "    标准RAG（检索增强生成）方法，基于Top-K检索。\n",
    "\n",
    "    Args:\n",
    "        pdf_path (str): 文档路径\n",
    "        query (str): 用户查询\n",
    "        k (int): 需要检索的块数量\n",
    "        chunk_size (int): 每个块的大小\n",
    "\n",
    "    Returns:\n",
    "        Dict: 包含查询、检索到的块和响应的结果字典\n",
    "    \"\"\"\n",
    "    print(\"\\n=== 开始标准TOP-K检索 ===\")\n",
    "    print(f\"查询: {query}\")\n",
    "\n",
    "    # 处理文档以提取文本、分块并创建嵌入\n",
    "    chunks, vector_store, doc_info = process_document(pdf_path, chunk_size)\n",
    "\n",
    "    # 为查询创建嵌入\n",
    "    print(\"创建查询嵌入并检索块...\")\n",
    "    query_embedding = create_embeddings([query])[0]\n",
    "\n",
    "    # 基于查询嵌入检索最相关的前k个块\n",
    "    results = vector_store.search(query_embedding, top_k=k)\n",
    "    retrieved_chunks = [result[\"document\"] for result in results]\n",
    "\n",
    "    # 将检索到的块格式化为上下文字符串\n",
    "    # context = \"\\n\\n\".join([\n",
    "    #     f\"CHUNK {i+1}:\\n{chunk}\"\n",
    "    #     for i, chunk in enumerate(retrieved_chunks)\n",
    "    # ])\n",
    "    context = \"\\n\\n\".join([\n",
    "        f\"文本块 {i+1}:\\n{chunk}\"\n",
    "        for i, chunk in enumerate(retrieved_chunks)\n",
    "    ])\n",
    "\n",
    "    # 使用上下文从语言模型生成响应\n",
    "    response = generate_response(query, context)\n",
    "\n",
    "    # 将结果编译成字典\n",
    "    result = {\n",
    "        \"query\": query,  # 用户查询\n",
    "        \"chunks\": retrieved_chunks,  # 检索到的块\n",
    "        \"response\": response  # 模型生成的响应\n",
    "    }\n",
    "\n",
    "    print(\"\\n=== 最终响应 ===\")\n",
    "    print(response)\n",
    "\n",
    "    return result\n"
   ],
   "id": "ee9aee391c6ca7cc",
   "outputs": [],
   "execution_count": 14
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 评估 RSE",
   "id": "9c572c243492287a"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-24T11:22:41.189379Z",
     "start_time": "2025-04-24T11:22:41.184370Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def evaluate_methods(pdf_path, query, reference_answer=None):\n",
    "    \"\"\"\n",
    "    比较RSE（相关段落提取）与标准Top-K检索方法。\n",
    "\n",
    "    Args:\n",
    "        pdf_path (str): 文档路径\n",
    "        query (str): 用户查询\n",
    "        reference_answer (str, 可选): 用于评估的参考答案\n",
    "    \"\"\"\n",
    "    print(\"\\n========= 评估开始 =========\\n\")\n",
    "\n",
    "    # 运行带有RSE（相关段落提取）的RAG方法\n",
    "    rse_result = rag_with_rse(pdf_path, query)\n",
    "\n",
    "    # 运行标准Top-K检索方法\n",
    "    standard_result = standard_top_k_retrieval(pdf_path, query)\n",
    "\n",
    "    # 如果提供了参考答案，则评估响应结果\n",
    "    if reference_answer:\n",
    "        print(\"\\n=== 比较结果 ===\")\n",
    "\n",
    "        # 创建一个评估提示，将两种响应与参考答案进行比较\n",
    "        evaluation_prompt = f\"\"\"\n",
    "            查询：{query}\n",
    "\n",
    "            参考答案：\n",
    "            {reference_answer}\n",
    "\n",
    "            标准检索系统回答内容：\n",
    "            {standard_result[\"response\"]}\n",
    "\n",
    "            相关片段提取系统回答内容：\n",
    "            {rse_result[\"response\"]}\n",
    "\n",
    "            请基于以下维度对比两个系统的回答内容质量：\n",
    "            1. 准确性与完整性\n",
    "            2. 问题解决有效性\n",
    "            3. 信息冗余度\n",
    "\n",
    "            具体要求：\n",
    "            • 判断哪个系统在各维度表现更优\n",
    "            • 对比分析时须直接引用具体内容\n",
    "            • 请逐项说明判断依据\n",
    "        \"\"\"\n",
    "\n",
    "        print(\"正在根据参考答案评估回答...\")\n",
    "\n",
    "        # 使用指定模型生成评估结果\n",
    "        evaluation = client.chat.completions.create(\n",
    "            model=llm_model,\n",
    "            messages=[\n",
    "                {\"role\": \"system\", \"content\": \"您是RAG系统回答内容的客观评估专家。\"},\n",
    "                {\"role\": \"user\", \"content\": evaluation_prompt}\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        # 打印评估结果\n",
    "        print(\"\\n=== 评估结果 ===\")\n",
    "        print(evaluation.choices[0].message.content)\n",
    "\n",
    "    # 返回两种方法的结果\n",
    "    return {\n",
    "        \"rse_result\": rse_result,\n",
    "        \"standard_result\": standard_result\n",
    "    }\n"
   ],
   "id": "9ab50cbf4c94f4e3",
   "outputs": [],
   "execution_count": 15
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-24T11:24:03.347564Z",
     "start_time": "2025-04-24T11:22:49.214213Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# Load the validation data from a JSON file\n",
    "with open('../../data/val.json', 'r', encoding=\"utf-8\") as f:\n",
    "    data = json.load(f)\n",
    "\n",
    "# Extract the first query from the validation data\n",
    "query = data[0]['question']\n",
    "\n",
    "# Extract the reference answer from the validation data\n",
    "reference_answer = data[0]['ideal_answer']\n",
    "\n",
    "# Run evaluation\n",
    "results = evaluate_methods(pdf_path, query, reference_answer)"
   ],
   "id": "9be1fe3b8bfb3bc",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "========= 评估开始 =========\n",
      "\n",
      "\n",
      "=== 开始带有相关段落提取的RAG ===\n",
      "查询: 什么是‘可解释人工智能’，为什么它被认为很重要？\n",
      "从文档中提取文本...\n",
      "将文本切分为非重叠段落...\n",
      "创建了 13 个块\n",
      "为块生成嵌入向量...\n",
      "\n",
      "计算相关性分数和块值...\n",
      "寻找最佳连续文本段落...\n",
      "找到段落 (0, 13)，得分 7.6610\n",
      "\n",
      "从块中重建文本段落...\n",
      "正在使用相关段落作为上下文生成响应...\n",
      "\n",
      "=== 最终响应 ===\n",
      "可解释人工智能（XAI，Explainable Artificial Intelligence）是指通过技术手段使人工智能系统的决策过程变得透明、可理解和可追溯的领域。其核心目标是解决传统AI模型（尤其是深度学习等复杂模型）的\"黑箱\"问题，让人类能够理解AI如何得出特定结论或做出特定决策。\n",
      "\n",
      "重要性体现在以下关键方面：\n",
      "\n",
      "1. 信任建立（第20章）\n",
      "- 透明度是建立用户对AI系统信任的基础\n",
      "- 可解释性让用户能够评估AI决策的可靠性和公平性\n",
      "\n",
      "2. 伦理合规（第19章）\n",
      "- 满足道德AI原则中的透明度和问责要求\n",
      "- 支持\"负责任AI\"的开发框架（第21章）\n",
      "- 是AI治理和监管的核心要求（第18章）\n",
      "\n",
      "3. 实际应用需求\n",
      "- 在医疗诊断（第11章）、金融风控（第7章）等高风险领域，决策依据必须可验证\n",
      "- 帮助发现和修正模型偏见（第4章提及的偏见放大问题）\n",
      "- 支持人机协作场景中的决策协调（第8章工作场景应用）\n",
      "\n",
      "4. 技术发展层面（第15章）\n",
      "- 是AI研究的前沿方向之一\n",
      "- 与\"以人为中心的AI\"发展理念深度契合\n",
      "- 促进更健壮、可靠的AI系统开发\n",
      "\n",
      "典型技术包括：\n",
      "- 决策可视化工具\n",
      "- 特征重要性分析\n",
      "- 局部可解释模型（LIME）\n",
      "- 反事实解释方法等\n",
      "\n",
      "随着AI在关键领域应用的深化（如第14章智慧城市、第12章网络安全），可解释性已成为确保AI系统安全、可靠、公平的必要条件，也是平衡技术创新与社会接受度的关键因素。\n",
      "\n",
      "=== 开始标准TOP-K检索 ===\n",
      "查询: 什么是‘可解释人工智能’，为什么它被认为很重要？\n",
      "从文档中提取文本...\n",
      "将文本切分为非重叠段落...\n",
      "创建了 13 个块\n",
      "为块生成嵌入向量...\n",
      "创建查询嵌入并检索块...\n",
      "正在使用相关段落作为上下文生成响应...\n",
      "\n",
      "=== 最终响应 ===\n",
      "**可解释人工智能（XAI）**是指通过技术手段使人工智能系统的决策过程透明化、易于理解的人工智能分支。其核心目标是破解AI模型的\"黑箱\"特性，让人类能够追溯和理解算法如何得出特定结论。\n",
      "\n",
      "**重要性主要体现在以下方面**：\n",
      "\n",
      "1. **信任建立**（文本块3/6/9）\n",
      "   - 当医疗诊断、金融风控等关键领域采用AI时，决策透明性直接影响用户信任度。XAI通过展示决策逻辑（如特征重要性分析、决策路径可视化）增强系统可信度。\n",
      "\n",
      "2. **伦理与责任**（文本块1/3/4）\n",
      "   - 涉及偏见的算法可能放大社会歧视（如招聘AI中的性别偏见）。XAI技术能识别偏差来源（文本块3），为算法审计提供依据，满足欧盟《AI法案》等合规要求。\n",
      "\n",
      "3. **安全验证**（文本块5/9）\n",
      "   - 在自动驾驶等高风险场景中，XAI可验证系统是否基于合理特征（如交通标志而非无关像素）做出判断，避免因数据噪声导致的致命错误。\n",
      "\n",
      "4. **性能优化**（文本块6）\n",
      "   - 可解释性分析能暴露模型弱点（如过度依赖某个非稳健特征），帮助开发者改进算法。深度学习领域特别关注通过注意力机制等XAI技术提升模型鲁棒性。\n",
      "\n",
      "5. **人机协作**（文本块1/7）\n",
      "   - 当AI作为人类决策辅助工具时（如医疗诊断），医生需要理解AI建议的依据以做出最终判断，XAI是实现有效协同的基础。\n",
      "\n",
      "当前主要技术包括LIME（局部解释）、SHAP值（特征贡献度量化）以及针对神经网络的层间可视化等。随着欧盟《数字服务法案》等法规强化AI透明度要求（文本块9），XAI已成为负责任AI开发的必备组件。\n",
      "\n",
      "=== 比较结果 ===\n",
      "正在根据参考答案评估回答...\n",
      "\n",
      "=== 评估结果 ===\n",
      "### 1. 准确性与完整性  \n",
      "**标准检索系统更优**  \n",
      "- **准确性**：标准检索系统提供了具体的技术示例（如LIME、SHAP值、注意力机制）和实际应用场景（医疗诊断、金融风控、自动驾驶），这些内容与XAI的核心定义和重要性高度吻合。例如，提到\"通过特征重要性分析、决策路径可视化增强信任\"（文本块3/6/9），直接关联XAI的技术实现。  \n",
      "- **完整性**：覆盖了信任、伦理、安全、性能优化、人机协作五大维度，并引用法规（如欧盟《AI法案》）和具体技术（层间可视化），信息全面。相比之下，相关片段系统的\"伦理合规\"部分仅泛泛提及\"道德AI原则\"（第19章），缺乏具体技术或案例支撑。  \n",
      "\n",
      "**引用对比**：  \n",
      "- 标准检索系统：*\"XAI技术能识别偏差来源（文本块3），为算法审计提供依据，满足欧盟《AI法案》等合规要求\"*（具体法规和技术结合）。  \n",
      "- 相关片段系统：*\"满足道德AI原则中的透明度和问责要求\"*（未说明如何满足）。  \n",
      "\n",
      "---\n",
      "\n",
      "### 2. 问题解决有效性  \n",
      "**标准检索系统更优**  \n",
      "- **针对性**：标准检索系统直接回答\"为什么重要\"，每个重要性论点均配以实际应用场景和技术手段。例如，针对\"安全验证\"提到\"避免因数据噪声导致的致命错误\"（文本块5/9），问题解决路径清晰。  \n",
      "- **实用性**：相关片段系统虽列出四大重要性（信任、伦理、应用需求、技术发展），但部分内容重复（如\"伦理合规\"与\"负责任AI\"框架），且未明确技术如何解决问题。例如，第4章提及\"偏见放大问题\"但未说明XAI如何修正。  \n",
      "\n",
      "**引用对比**：  \n",
      "- 标准检索系统：*\"可解释性分析能暴露模型弱点（如过度依赖某个非稳健特征），帮助开发者改进算法\"*（具体问题+解决方案）。  \n",
      "- 相关片段系统：*\"帮助发现和修正模型偏见\"*（未说明修正方法）。  \n",
      "\n",
      "---\n",
      "\n",
      "### 3. 信息冗余度  \n",
      "**相关片段系统更优**  \n",
      "- **简洁性**：相关片段系统通过章节引用（如第11章医疗诊断、第7章金融风控）概括领域应用，避免了过多技术细节，冗余度低。标准检索系统虽信息全面，但部分内容重复（如\"信任建立\"在文本块3/6/9中多次出现）。  \n",
      "- **结构化**：相关片段系统用分点（1. 信任建立，2. 伦理合规等）清晰划分重要性，而标准检索系统的\"性能优化\"与\"安全验证\"部分内容有重叠（均涉及模型可靠性）。  \n",
      "\n",
      "**引用对比**：  \n",
      "- 相关片段系统：*\"典型技术包括：决策可视化工具、特征重要性分析\"*（仅列举技术名称）。  \n",
      "- 标准检索系统：*\"当前主要技术包括LIME（局部解释）、SHAP值（特征贡献度量化）以及针对神经网络的层间可视化等\"*（详细解释技术，可能超出必要信息）。  \n",
      "\n",
      "---\n",
      "\n",
      "### 总结  \n",
      "- **标准检索系统**在**准确性**和**问题解决有效性**上显著胜出，因其提供具体技术、法规和场景化解释。  \n",
      "- **相关片段系统**在**信息冗余度**上更优，结构简洁但牺牲了部分深度。  \n",
      "- 若用户需要专业解答（如技术人员），推荐标准检索系统；若需快速概览，相关片段系统更合适。\n"
     ]
    }
   ],
   "execution_count": 16
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
